#' @title immunetherapy_FUN 子程序，根据coef计算得分
#' @description 根据输入建模方式或者自定义建模方式，来进行计算得分,也可以自定义function调用exprs、x或者其他输入参数计算得分。
#' @param model_coef 只有两列的data.frame，第一列为基因，第二列为coef
#' @param exprs 表达谱
#' @export 
#' @author *WYK*
#'
compute_score_in_icb <- function(exprs, model_coef = NULL) {
  exprs <- exprs[model_coef[, 1], , drop = F] %>%
    t() %>%
    as.data.frame()

  miss_gene <- model_coef[, 1][which(is.na(exprs[1, ]))]
  miss_gene_symbol <- paste0(miss_gene, collapse = ", ")

  exprs[is.na(exprs)] <- 0

  riskscore_res <- as.matrix(exprs) %*% as.matrix(model_coef[, 2, drop = F]) %>%
    as.data.frame() %>%
    dplyr::rename(riskscore = 1)

  if (length(miss_gene) == 1) {
    cli::cli_alert_info(str_glue("'WARNING:' {miss_gene_symbol} 在当前队列的表达谱中没有匹配到。\n"))
  } else {
    cli::cli_alert_info(sprintf("所用基因在当前队列表达谱中没有丢失"))
  }

  return(riskscore_res)
}

#' @title immunetherapy_FUN 子程序，根据就不同方法计算得分
#' @description 根据输入建模方式或者自定义建模方式，来进行计算得分,也可以自定义function调用exprs、x或者其他参数计算得分。
#' @param model 字符串，已有的建模方法lasso,multicox,PCA1,PCA2。可以选择其中一个进行分析，也可以自定义
#' @param x 根据model中不同的方式，给到不同的东西，比如说lasso方法需要给到基因coef相关的data.frame，PCA1需要的是基因字符串向量，PCA2需要给到一个list等等
#' @param exprs 表达谱
#' @param ... 可以自定义建模方法。把`exprs`与`x`或者自定义的其他参数，传递到自定义函数内部的function，计算riskscore。自定义函数可以是匿名函数
#' 或者直接给一个riskscore，格式有要求：样本在行名，且只有一列，第一列表头为riskscore，内容为各个样本的得分。
#' @return 只有一列的data.frame，行名为样本名称，列名为riskscore
#' @usage
#' score_res <- get_score(
#'  model = "means",
#'  x = sample(rownames(train_data$tumor_exprs)),
#'  exprs = train_data$tumor_exprs,
#'  function() {
#'    exprs[x, ] %>%
#'      t() %>%
#'      as.data.frame() %>%
#'      rowMeans() %>%
#'      as.data.frame() %>%
#'      rename(riskscore = 1)
#'  }
#' )
#' @export
#' @author *WYK*
#'
get_score <- function(model, x, exprs, ...) {
  if (!missing(...) && !{
    model %in% c("lasso", "multicox", "PCA1_minus_PCA2", "PCA1_sum_PCA2",'ssgsea')
  }) {
    fun_defined <- \(...) {}
    lst_inner <- list(...)
    for (i in names(lst_inner)) {
      if (i != "") {
        assign(i, lst_inner[[i]])
      }
    }
    body(fun_defined) <- body(lst_inner[[which(map_lgl(lst_inner, ~ is.function(.x)))]])
    score_res <- fun_defined()
    return(score_res)
  }

  signature_coef
  model_fun <- list(
        lasso = \(...){
            if (!is.data.frame(x)) {
                cli::cli_alert_danger(str_glue("{model}要求输入x是DF，且第一列为基因名，第二列为相关系数"))
                return()
            }
            score_res <- compute_score_in_icb(exprs, model_coef = x)
            return(score_res)
        },
        multi_cox = \(...){
            if (!is.data.frame(x)) {
                cli::cli_alert_danger(str_glue("{model}要求输入x是DF，且第一列为基因名，第二列为相关系数"))
                return()
            }
            score_res <- compute_score_in_icb(exprs, model_coef = x)
            return(score_res)
        },
        PCA1_minus_PCA2 = \(...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("x中不存在{.var gene}列")
            }
            if (!any(grepl(pattern = "HR", x = colnames(x)))) {
                cli::cli_alert_danger("x中不存在{.var HR}列")
            }

            up_gene <- x %>%
                filter(HR > 1) %>%
                pull(gene)
            down_gene <- x %>%
                filter(HR < 1) %>%
                pull(gene)

            x_lst <- vector("list", 2)
            pacman::p_load(FactoMineR)

            x_lst[[1]] <- intersect(up_gene, rownames(exprs))
            x_lst[[2]] <- intersect(down_gene, rownames(exprs))

            if (length(x_lst[[1]]) <= 2 || length(x_lst[[2]]) <= 2 ) {
                cli::cli_alert_danger("在方法{.var PCA1_minus_PCA2}中，输入的基因列表的第一个或第二个元素在表达谱中交集小于2个，程序跳出。")
                return()
            }

            pcaH <- PCA(X = exprs[x_lst[[1]], ] %>% t(), graph = FALSE, ncp = 2)
            pcaL <- PCA(X = exprs[x_lst[[2]], ] %>% t(), graph = FALSE, ncp = 2)
            score <- pcaH$ind$contrib[, 1] + pcaH$ind$contrib[, 2] - pcaL$ind$contrib[, 1] - pcaL$ind$contrib[, 2]
            pca_res <- data.frame(PC1PC2 = score) %>% t()

            rownames(pca_res) <- 'Score'

            score_res <- compute_score_in_icb(exprs = pca_res, data.frame(gene = "Score", coef = 1))
            return(score_res)
        },
        PCA1_sum_PCA2 = \(...) {
            if (!any(grepl(pattern = "gene", x = colnames(x)))) {
                cli::cli_alert_danger("单因素cox结果中不存在{.var gene}列")
                return()
            }

            x <- x[[1]]

            x <- intersect(x, rownames(exprs))
            pca_res <- prcomp(x = exprs[x, ], retx = F, scale = T, center = F)

            pca_res <- pca_res$rotation %>%
                as.data.frame() %>%
                transmute(PC1PC2 = PC1 + PC2) %>%
                rename(Score = 1) %>%
                t() %>%
                as.data.frame()

            score_res <- compute_score_in_icb(exprs = pca_res, data.frame(gene = "Score", coef = 1))
            return(score_res)
        },
        ssgsea = \(...) {
            f <- function(...) {
                d <- GSVA::gsva(
                    expr = exprs %>% na.omit() %>% as.matrix(),
                    gset.idx.list = list("Score" = x[[1]]),
                    method = "ssgsea",
                    kcdf = "Gaussian", abs.ranking = FALSE,
                    min.sz = 2,
                    ssgsea.norm = T,
                    verbose = T,
                    parallel.sz = 4
                ) %>%
                    as.data.frame() 
                return(d)
            }

            ssgsea_res <- f()

            score_res <- compute_score_in_icb(exprs = ssgsea_res, data.frame(gene = "Score", coef = 1))
            return(score_res)
        }
    )

  return(model_fun[[model]]())
}


#' @title immunetherapy_FUN 子程序，根据之前算好的score，运行KM分析
#' @description 本质上就是利用多种方法进行KM分析，调用ComplexKM_v2进行分析
#' @return list 第一个为KM曲线gg对象，第二个为计算好的df，包含score等信息
#' @export 
#' @author *WYK*
#'
do_KM <- function(...) {
  KM_method_in_ComplexKM <<- match.arg(KM_method,
    c("BestCut", "NormalCut", "TertileCut", "QuartileCut"),
    several.ok = TRUE
  )

  km_method_index <- as.list(rep(FALSE,4))
  names(km_method_index) <- c("BestCut", "NormalCut", "TertileCut", "QuartileCut")
  index <- map_int(KM_method_in_ComplexKM, ~ match(.x, names(km_method_index)))
  km_method_index[index] <- TRUE

  Ccomplexkm_res <- ComplexKM_v2(
    Input = clinical_for_km,
    surtime_unit = surtime_unit,
    time_col = time_col,
    status_col = status_col,
    score_col = riskscore_name,
    BestCut = km_method_index$BestCut,
    NormalCut = km_method_index$NormalCut,
    TertileCut = km_method_index$TertileCut,
    QuartileCut = km_method_index$QuartileCut,
    CutOffManual = NULL,
    GroupManul = NULL,
    legend = "top", SaveFile = F,
    od = od, w = 5, h = 5
  )

  KM_method_in_loop <- match.arg(KM_method_in_ComplexKM,
    c("BestCut_KM", "NormalCut_KM", "TertileCut_KM", "QuartileCut_KM"),
    several.ok = TRUE
  )

  km_res_filtered <- map(KM_method_in_loop, function(x) {
    km_p <- Ccomplexkm_res[[x]][[1]]
    km_group <- Ccomplexkm_res[[x]][[3]]
    return(list("KM_p" = km_p, "KM_group" = km_group))
  })

  names(km_res_filtered) <- KM_method_in_ComplexKM

  return(km_res_filtered)
}


#' @title 根据不同建模方法预测免疫治疗结果
#' @description 本function灵活考虑了多种建模方法，可以自定义riskscore计算方法
#' @inheritParams get_score
#' @param clin 临床信息df
#' @param time_col 临床信息df中 时间所在列
#' @param status_col 临床信息df中，状态所在列
#' @param riskscore_name 得分名字
#' @param KM_method KM计算方法，详情见ComplexKM_v2，可以是c("BestCut_KM", "NormalCut_KM", "TertileCut_KM", "QuartileCut_KM")其中一个两个，
#' 也可以模糊匹配，比如说可以用c('N','B')代表其中的中位数分割方法，以及最优分组分割方法
#' @param color_in_plot 除KM外的配色
#' @param od 结果输出路径
#' @param w 图宽
#' @param h 图高
#' @param file_name 字符串，用于命名文件以及文件夹的字段
#' @usage
#' load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/data/forest_data.RData")
#' load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/out/2.major_part/major.RData", verbose = T)
#' load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/data/train_data.RData", verbose = T)
#' test <- immunetherapy_FUN(
#'  model = "means",
#'  x = sample(rownames(train_data$tumor_exprs)),
#'  exprs = train_data$tumor_exprs,
#'  function() {
#'    exprs[x, ] %>%
#'      t() %>%
#'      as.data.frame() %>%
#'      rowMeans() %>%
#'      as.data.frame() %>%
#'      rename(riskscore = 1)
#'  },
#'  clin = forest_data %>% mutate(Immunotherapy.Response.Group = ifelse(status ==1,'R','NoR')),
#'  time_col = "time",
#'  status_col = "status",
#'  riskscore_name = "ASS_SCORE",
#'  KM_method = c("Best",'N'),
#'  color_in_plot = NULL,
#'  savefile = T,
#'  od = "/Pub/Users/wangyk/project/Poroject/GAP220630002_LIHC_Age_gene_SC_bulk/test",
#'  w = 4,
#'  h = 4,
#'  file_name = "12344"
#' )
#'
#' @export
#' @return NULL
#' @author *WYK*
#'
immunetherapy_FUN <- function(model, x, exprs, ...,
                              clin, time_col = "time", status_col = "status", riskscore_name = "riskscore",
                              KM_method = c("Best_Cut"), color_in_plot = NULL, savefile = T, surtime_unit = 365,
                              od, w, h, file_name = "cohort") {
  # 调用脚本：
  source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/NewLover/ComplexKM.R")
  source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/NewLover/ClinicalFeatures_plot.R")

  pacman::p_load(rlang, magrittr, tidyverse, magrittr)

  coef_1_df <- data.frame(gene = "Score", coef = 1)

  # cli::cli_alert_info(crayon::bold(str_glue("\n\n当前分析：{file_name} >>>")))
  
  score_res <- get_score(model = model, x = x, exprs = exprs, ... = ...)
  #  score_res <- get_score(model , x, exprs)

  if (is.null(score_res)) {
    cli::cli_alert_danger("得分计算失败，跳过当前队列。")
    return()
  }

  if (!is.data.frame(score_res)) {
    cli::cli_alert_danger(text = "得分计算错误，不是data.frame格式，跳过当前队列。要求计算结果为，样本在行名，第一列列名为riskscore的data.frame。")
    return()
  }

  if (colnames(score_res)[1] != 'riskscore') {
    cli::cli_alert_danger("得分计算错误，跳过当前队列。\n要求计算结果为，样本在行名，第一列列名为riskscore的data.frame。")
    return()
  }

  if(sum(score_res[[1]]) == 0){
    cli::cli_alert_danger("得分计算错误，跳过当前队列。得分全部为0。")
    return()
  }

  score_res %<>% dplyr::select(riskscore) %>% dplyr::rename(!!riskscore_name := "riskscore")

  clinical_for_all_patients <- suppressMessages(clin %>% full_join(rownames_to_column(score_res, "Sample"))) 
  clinical_for_km <- suppressMessages(clin %>% inner_join(rownames_to_column(score_res, "Sample")))

  if (all(c(time_col, status_col) %in% colnames(clin))) {
    km_fun_inner <- \(...){}
    body(km_fun_inner) <- body(do_KM)
    km_res <- km_fun_inner()
  } else {
    cli::cli_alert_info("临床信息有缺失，KM分析跳过")
  }

  if ("Immunotherapy.Response.Group" %>% is_in(colnames(clin))) {
    clinical_for_all_patients_ <- clinical_for_all_patients %>% rename("Response " = "Immunotherapy.Response.Group")
    box_in_all_patients <- ClinicalFeatures_plot(
      Data = clinical_for_all_patients_, phenotype_col = "Response ", score_col = riskscore_name,
      color_used = color_in_plot, saveplot = F, alpha = .9
    )

    score_in_all_patient <- box_in_all_patients$Clinical_in_Score[["p_list"]][["Response "]]

    if (isTRUE(savefile)) {
      od_now <- str_glue("{od}/{file_name}/")
      plotout(
        od = od_now,
        name = str_glue("Score_in_all_patients_{file_name}"), w = w * .82, h = h,
        p = score_in_all_patient
      )
    }

    if (!exists("km_res", envir = current_env())) {
      return()
    }

    os_related_plots <- map(KM_method_in_ComplexKM, function(method) {
      df <- km_res[[method]][["KM_group"]]

      if ("Response" %in% colnames(df)) {
        df %<>% rename("Response_raw" = "Response")
      }

      df %<>% rename("Response" = "Immunotherapy.Response.Group") %>% rename(!!riskscore_name := "Score")

      box_in_os_patients <- ClinicalFeatures_plot(
        Data = df, phenotype_col = "Response", score_col = riskscore_name,
        group_col = "Group",
        color_used = color_in_plot, saveplot = F, alpha = .9
      )
      score_in_os_patients_response <- box_in_os_patients$Clinical_in_Score[["p_list"]][["Response"]]
      percentage_stack_in_os_patients <- box_in_os_patients[["Clinical_in_Group"]][["p_list_group"]][["Response"]]

      res <- list(
        "score_in_os_patients" = score_in_os_patients_response,
        "percentage_stack_in_os_patients" = percentage_stack_in_os_patients
      )

      return(res)
    })
    names(os_related_plots) <- KM_method_in_ComplexKM

    if (isTRUE(savefile)) {
      walk(KM_method_in_ComplexKM, function(method) {
        od_now <- str_glue("{od}/{file_name}/{method}/")

        plotout(
          od = od_now,
          name = str_glue("KM_{file_name}"), w = w, h = h,
          p = km_res[[method]][["KM_p"]]
        )

        plotout(
          od = od_now,
          name = str_glue("Score_in_response_{file_name}"), w = w * .82, h = h,
          p = os_related_plots[[method]][["score_in_os_patients"]]
        )

        plotout(
          od = od_now,
          name = str_glue("Percentage_Stack_{file_name}"), w = w * .82, h = h,
          p = os_related_plots[[method]][["percentage_stack_in_os_patients"]] + guides(fill = guide_legend(nrow = 2))
        )

        p_ <- cowplot::plot_grid(
          km_res[[method]][["KM_p"]],
          os_related_plots[[method]][["score_in_os_patients"]],
          os_related_plots[[method]][["percentage_stack_in_os_patients"]] + guides(fill = guide_legend(nrow = 2)),
          nrow = 1, labels = "AUTO",greedy = c(1,.9,.9)
        )

        plotout(
          od = od_now,
          name = str_glue("{file_name}"), w = w * 2.7, h = h,
          p = p_
        )

        write_tsv(
          x = km_res[[method]][["KM_group"]],
          file = str_glue("{od_now}/Table_Clinical_Score_{file_name}.tsv")
        )
      })
    }
  } else {
    cli::cli_alert_info(crayon::blue("该队列反应信息有缺失，Score是否有显著分布差异以及百分比堆积图等分析跳过"))
  }

  if (!exists("km_res", envir = current_env())) {
    return()
  }

  if (isTRUE(savefile)) {
    walk(KM_method_in_ComplexKM, function(method) {
      od_now <- str_glue("{od}/{file_name}/{method}/")

      plotout(
        od = od_now,
        name = str_glue("KM_{file_name}"), w = w, h = h,
        p = km_res[[method]][["KM_p"]]
      )
      write_tsv(
        x = km_res[[method]][["KM_group"]],
        file = str_glue("{od_now}/Table_Clinical_Score_{file_name}.tsv")
      )
    })
  }

  return(km_res[[length(km_res)]])

}

# # 测试部分------
# load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/data/forest_data.RData")
# load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/out/2.major_part/major.RData", verbose = T)
# load("/Pub/Users/wangyk/project/Poroject/GDT220812001_GBM_immuneCellDeath_WGCNA/data/train_data.RData", verbose = T)
# library(YK.pkg)
# test <- immunetherapy_FUN(
#   model = "means",
#   x = sample(rownames(train_data$tumor_exprs)),
#   exprs = train_data$tumor_exprs,
#   function() {
#     exprs[x, ] %>%
#       t() %>%
#       as.data.frame() %>%
#       rowMeans() %>%
#       as.data.frame() %>%
#       rename(riskscore = 1)
#   },
#   clin = forest_data %>% mutate(Immunotherapy.Response.Group = ifelse(status == 1, "R", "NoR")),
#   time_col = "time",
#   status_col = "status",
#   riskscore_name = "ASS_SCORE",
#   KM_method = c("Best", "N"),
#   color_in_plot = NULL,
#   savefile = T,
#   od = "/Pub/Users/wangyk/project/Poroject/GAP220630002_LIHC_Age_gene_SC_bulk/test",
#   w = 4,
#   h = 4,
#   file_name = "12344"
# )

# -----


# source("/Pub/Users/wangyk/Project_wangyk/Codelib_YK/some_scr/NewLover/immunetherapy_FUN.R")
# load("/Pub/Users/cuiye/projects/F230526001/results/6.2.caf_signature/signature_coef.RData")

# x = signature_coef
